clear;

addpath(genpath("./src"));

PLOT_SIM = true;
DATA_CAST = 'gpuArray-single';

nx = 256;
dx = 0.1e-3;
ny = 256;
dy = 0.1e-3;

nt = 4096;
fs = 250e6;
fs_div = 4;
fc = 5e6;
pitch = 0.6e-3;
width = 0.5e-3;
num_elem = 32;
pulse_width = 4;
c_vec = [1483, 5900];

array_angle = 0;
array_pos = [0, -10e-3];

kgrid = kWaveGrid(nx, dx, ny, dy);
kgrid.t_array = (0: (nt-1)) * (1 / fs);

cmap = ones(ny, nx) * c_vec(2);
% cmap(:, kgrid.y_vec > 0) = c_vec(2);

flaw_x = [-8e-3, 0, 8e-3];
flaw_y = [3e-3, 6e-3, 9e-3];

for i = 1 : numel(flaw_x)
    [~, tmp_x] = min(abs(kgrid.x_vec - flaw_x(i)));
    [~, tmp_y] = min(abs(kgrid.y_vec - flaw_y(i)));

    tmp = makeDisc(nx, ny, tmp_x, tmp_y, 2);
    cmap(tmp == 1) = 343;
end

[array_mask, array_mask_idx, elem_pos] = funcGenLinearArrayMask(num_elem, pitch, width, array_angle, array_pos, kgrid);

sensor.mask = array_mask;
source.p = toneBurst(fs, fc, pulse_width);

medium.sound_speed = cmap;
medium.density = ones(nx, ny) * 1e3;

%% stimulate all elements
source.p_mask = array_mask;

input_args = {'PlotSim', PLOT_SIM, 'PlotPML', false, 'PlotLayout', false, 'DisplayMask', source.p_mask, ...
    'PMLAlpha', 2, 'PMLInside', true, 'DataCast', DATA_CAST, 'DataRecast', true, 'RecordMovie', false};
sensor_data = kspaceFirstOrder2D(kgrid, medium, source, sensor, input_args{:});

%% stimulate a single element
rfdata = zeros(nt / fs_div, num_elem);
for i = 1 : num_elem
    source.p_mask = zeros(nx, ny);
    source.p_mask(array_mask_idx == i) = 1;

    sensor.mask = source.p_mask;
    
    input_args = {'PlotSim', false, 'PlotPML', false, 'PlotLayout', false, 'DisplayMask', source.p_mask, ...
        'PMLAlpha', 2, 'PMLInside', true, 'DataCast', DATA_CAST, 'DataRecast', true, 'RecordMovie', false};
    sensor_data = kspaceFirstOrder2D(kgrid, medium, source, sensor, input_args{:});

    rfdata(:, i) = sum(sensor_data(:, 1 : fs_div:end), 1);
end

%% stimulate a single element and receive signals of all elements
rfdata = zeros(nt/fs_div, num_elem, num_elem);

sensor.mask = array_mask;
array_mask_vec = array_mask_idx(array_mask_idx ~= 0);
for i = 1 : 2
    source.p_mask = zeros(nx, ny);
    source.p_mask(array_mask_idx == i) = 1;

    input_args = {'PlotSim', PLOT_SIM, 'PlotPML', false, 'PlotLayout', false, 'DisplayMask', source.p_mask, ...
        'PMLAlpha', 2, 'PMLInside', true, 'DataCast', DATA_CAST, 'DataRecast', true, 'RecordMovie', false};
    sensor_data = kspaceFirstOrder2D(kgrid, medium, source, sensor, input_args{:});

    for j = 1 : num_elem
        rfdata(:, j, i) = sum(sensor_data(array_mask_vec == j, 1:fs_div:end), 1);
    end
end

%%
img = das(hilbert(rfdata), kgrid.t_array(1:fs_div:end), elem_pos, kgrid.x_vec, kgrid.y_vec, c_vec(2));

figure(1);
imagesc(kgrid.x_vec, kgrid.y_vec, abs(img(:, :, 1)));
axis equal tight;

figure(2);
imagesc(kgrid.x_vec, kgrid.y_vec, abs(sum(img, 3)));
axis equal tight;

%% plots
figure(1);
imagesc(kgrid.y_vec, kgrid.x_vec, cmap);
axis equal tight;
colorbar;
xlabel("y [m]");
ylabel("x [m]");

figure(2);
imagesc(kgrid.x_vec, kgrid.y_vec, array_mask);
axis equal tight;
colorbar;

figure(3);
imagesc(kgrid.x_vec, kgrid.y_vec, array_mask_idx);
axis equal tight;
colorbar;

figure(4);
plot(source.p);
